NanoVDB: chain DeviceBuffer::recordUse and expose orderAfterPriorUses (CUDA) - #2287
Conversation
… (CUDA)
recordUse re-records the single per-device tracking event, and re-recording
moves an event: uses recorded on streams A then B left only B's coverage, so
the device free could run while A's work was in flight — violating the
invariant documented on mEvents ("every use waits on this event before
issuing work and re-records it afterwards"). Record now waits on the prior
capture first, so the event transitively covers every recorded use. Also make
orderAfterPriorUses public as the consume-side companion, so external
consumers (e.g. the Python bindings' zero-copy exports) can order their own
stream after the tracked uses. Adds a DeviceBufferChainedRecordUse regression
test that fails before the fix and passes after.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Jonathan Swartz <jonathan@jswartz.info>
harrism
left a comment
harrism
left a comment
There was a problem hiding this comment.
Solid fix — the event-move hazard is real, the chaining mechanism is correct (verified the never-recorded-event no-op, cross-device wait legality, first-creation branch, and per-device chain independence), and the A/B-verified regression test is exactly the right discipline. Three comments inline: one design trade worth acknowledging (concurrent-reader serialization), one test-robustness suggestion (vacuous pass when the block isn't recycled), and an optional style nit.
| } else { | ||
| // Re-recording MOVES the event; chain first so the new capture also covers the | ||
| // prior recorded use (waiting on a never-recorded or completed event is a no-op). | ||
| cudaCheck(cudaStreamWaitEvent(stream, mEvents[device], 0)); |
There was a problem hiding this comment.
The fix is right for the motivating case (write on A, then record on B), but note it also changes recordUse from an observer into an ordering mutator for concurrent readers: two readers recording uses on streams A then B were previously independent, and now B serializes behind A's capture point. The single-event design can't distinguish read‖read from write→read, and the doc note's "later consumers observing earlier writes is the expected ordering" only covers the latter. Correctness over concurrency is the right trade for a tracking convenience — but do the #2225 bindings (or other consumers) fan out concurrent readers today? If so it might deserve a sentence in the doc note; if a profile ever surfaces this, the fix would be a read/write-separated or per-record event scheme rather than a revert.
There was a problem hiding this comment.
Agreed on the trade, and to your question: no — the #2225 bindings never call recordUse internally. The CAI/DLPack exports call orderAfterPriorUses (the consume-side wait), which doesn't record and so doesn't serialize anything; recordUse is only ever user-invoked from Python. So the read‖read serialization is strictly opt-in today, and the cost is one event-wait on the recording stream. I've extended the doc note in 5237e3e to state the concurrent-reader serialization explicitly and name the upgrade path (read/write-separated or per-record events, not a revert).
| cudaCheck(cudaStreamDestroy(userB)); | ||
| cudaCheck(cudaStreamDestroy(other)); | ||
|
|
||
| EXPECT_EQ(0u, clobbered) << "a later recordUse on another stream discarded the tracking " |
There was a problem hiding this comment.
The detection here depends on the freed block being recycled into victim; when the allocator doesn't recycle (different pool state, or a device without pool support), clobbered is 0 and the test passes without having tested the chain. It reports recycled in the failure message but never requires it — so silence is indistinguishable from success. Inherited from the FreeOrdering harness family, so arguably out of scope here, but a one-liner like if (!recycled) GTEST_SKIP() << "allocator did not recycle the block; ordering not exercised"; would make a vacuous run visible instead of green.
There was a problem hiding this comment.
Adopted in 5237e3e, with one check first: I wanted to be sure non-recycling couldn't be a consequence of the fix (in which case the skip would fire precisely on success). Probed on sm_120 across repeated trials post-fix: recycled=1, stillPending=1 every time — the stream-ordered pool hands the same block back with the dependency attached rather than refusing to recycle, so on a working platform the skip never fires and the clobber detection is genuinely exercised. On a non-recycling platform it converts a silent vacuous green into a visible SKIP. I scoped it to the new test; happy to apply the same guard to the two inherited FreeOrdering tests as a follow-up if you'd like.
| /// partially-written buffer after asynchronous uploads or recorded kernels. | ||
| /// @param device Device whose buffer is about to be read | ||
| /// @param stream Stream the consumer's work will be issued on | ||
| void orderAfterPriorUses(int device, cudaStream_t stream) const |
There was a problem hiding this comment.
Optional style nit: this publishes the method by carving a public:/private: island inside the private section. Moving the method down to the public block keeps the access sections contiguous, which is the pattern elsewhere in the file.
There was a problem hiding this comment.
Done in 5237e3e — moved into the public section directly above recordUse, so the produce/consume halves of the tracking API sit together and the access sections stay contiguous.
…ass, tidy access sections - recordUse's doc note now states that concurrent READERS recording uses serialize behind each other (the single per-device event cannot distinguish read-read from write-read) and names the upgrade path (read/write-separated or per-record events) should a profile ever surface it. - DeviceBufferChainedRecordUse now GTEST_SKIPs when the allocator did not recycle the freed block, making a vacuous run visible instead of green. Verified on sm_120 that stream-ordered pools recycle WITH the dependency attached even post-fix, so the skip does not fire on a working platform. - orderAfterPriorUses moved into the public section next to recordUse instead of a public/private island inside the private section. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Jonathan Swartz <jonathan@jswartz.info>
kmuseth
left a comment
kmuseth
left a comment
There was a problem hiding this comment.
Approved as-is
Here is Calude's summary:
The bug: cuda::DeviceBuffer tracks one CUDA event per device to know when it's safe to free the buffer. cudaEventRecord moves an event to a new capture point rather than adding to it. So if recordUse was called on stream A and then stream B, B's record silently discarded the only tracking of A's in-flight work — the buffer's destructor/clear()/move-assignment could then free memory while A was still using it, and CUDA's stream-ordered allocator could hand that memory to a new, unrelated allocation before A's write landed. This is a real use-after-free / data-corruption bug, not just a style issue.
The fix: recordUse now does cudaStreamWaitEvent(stream, event) before re-recording, so the new capture point is only reached after the previous capture point is reached too — the single event becomes a proper chain covering every recorded use transitively, matching the invariant already documented on mEvents and already followed by the internal deviceUpload/deviceDownload paths. orderAfterPriorUses (the consumer-side wait, previously private) is made public so external code (e.g. the Python array-interface/DLPack export work referenced in the PR) can order its own reads correctly.
Verification performed:
Traced the ordering logic by hand: since cudaStreamWaitEvent only affects work enqueued after it on that stream, and callers enqueue their actual buffer-use work before calling recordUse, the chain correctly orders "new capture" after both "this stream's own use" and "whatever was tracked before" — without blocking the host or the already-enqueued work.
Confirmed the honestly-documented side effect (subsequent work on the recording stream also waits on the prior use, and concurrent read-only uses get serialized since one event can't distinguish read-read from write-read) is an acceptable, disclosed tradeoff rather than a hidden regression.
Checked freeDeviceBuffers, deviceUpload, and deviceDownload — confirmed they already called orderAfterPriorUses/recordUse correctly, and that recordUse's new internal wait is at worst a harmless redundant no-op there (waiting again on an event that hasn't changed since the caller's own explicit wait moments earlier).
Built and ran the new DeviceBufferChainedRecordUse test on real hardware (RTX 6000 Ada) — compiled the full 4300-line TestNanoVDB.cu cleanly, linked, and ran it.
Did the A/B check myself, not just trusted the PR's claim: reverted only the recordUse chaining fix in a worktree, rebuilt, and ran the new test 5 times — it failed deterministically all 5 times, with the exact symptom the PR describes (block recycled: true, work still pending when it was reused: true — the pool handed the freed block to a new allocation while the original write was still in flight). Ran the fixed version 5 times — passed all 5.
Ran the full TestNanoVDBCUDA suite (61 tests) with the fix applied: 60 passed, and the one failure (UnifiedBuffer_IO, missing a test-data fixture file) is precisely the pre-existing, unrelated exception the PR description called out — confirmed independently, not something I had to take on faith.
Ran the new test plus the related free-ordering tests under compute-sanitizer --tool racecheck: 0 hazards.
Findings: none. This is a correct, well-isolated fix for a real race, backed by a test that I confirmed genuinely fails without the fix and passes with it — about as strong a verification as this kind of change gets.
Summary
Two small, coupled changes to
cuda::DeviceBuffer's use tracking, prompted by review discussion on #2225'srecordUsebindings.1.
recordUsenow chains across streams (bug fix). The buffer tracks one event per device, andcudaEventRecordmoves an event. So with uses recorded on streams A then B, B's record discarded the only coverage of A's in-flight work, and the device free (destructor /clear()/ move-assignment) could run while A was still using the buffer. This violated the invariant already documented onmEvents("every use waits on this event before issuing work and re-records it afterwards") — the internal upload/download paths follow it, but the publicrecordUseskipped the wait half.recordUsenow issuescudaStreamWaitEvent(stream, event)before re-recording, so the single event transitively covers every recorded use. Side effect (documented): work subsequently issued on the recording stream is also ordered after the previously recorded use — for a shared buffer that is the expected visibility ordering.2.
orderAfterPriorUsesis now public (API). It is the consume-side companion ofrecordUse: an external consumer (e.g. the Python bindings'__cuda_array_interface__/ DLPack exports in #2225) calls it with its own stream before reading, so it cannot observe a partially-written buffer after an asynchronous upload. It was previously private, leaving no public way to order a consumer stream against the tracked uses.Testing
TestNanoVDBCUDA.DeviceBufferChainedRecordUseregression test, modeled on the existingDeviceBuffer*FreeOrderingharness: a non-blocking stream is parked behind a busy-wait with a recorded late write, a second idle stream records a use after it, and the buffer is destroyed. Verified A/B: fails before the fix ("block recycled: true", the recycled allocation is clobbered by the late write) and passes after.nanovdb_test_cudasuite passes on a Blackwell (sm_120) GPU (the one pre-existing exception,UnifiedBuffer_IO, requires a fixture written by the host suite'sGridCountAndIndexand passes when run after it — unrelated to this change).🤖 Generated with Claude Code